#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sub license,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
# Mountain View, CA
#

""" assemble pde matrices """

from scipy import dot, diag, c_, zeros


class AssembleStiffnesMassMatrices(object):

    """
      Assemble mass and stiffness matrices.
      The coefficients can be constant or given by splines.
      Quadrature formulas are needed for numerical integration
      ( Spline coefficients )
      NOTE:
        Assembling matrices for constant coefficient PDE's is trivial
        and can be done much more efficiently than what is done here.
        This class tries to unify the concept of 'assembling' in general.
        It encapsulate the fact that once you know the behavior at the
        element level, the rest is just a matter of gluing them together.
    """

    __slots__ = ["m_grid", "m_basis", "m_quads", "m_c", "m_d", "m_r",
                 "m_mass_diag", "m_reac_diag", "m_stiff_diag", "m_reac",
                 "m_conv", "m_diff", "m_c_is_real", "m_d_is_real",
                 "m_r_is_real", "m_der_at_quad", "m_shape_at_quad", "m_zero"]

    def __init__(self, grid):
        self.m_grid = grid
        (self.m_basis, self.m_quads) = (None, None)

        (self.m_c, self.m_d, self.m_r) = (None, None, None)
        (self.m_reac, self.m_conv, self.m_diff) = (None, None, None)

        (self.m_mass_diag, self.m_reac_diag, self.m_stiff_diag) = \
            (None, None, None)

        (self.m_c_is_real, self.m_d_is_real, self.m_r_is_real) = \
            (True, True, True)

        (self.m_der_at_quad, self.m_shape_at_quad) = (None, None)
        self.m_zero = None

    def set_new_grid(self, grid):
        """ set new grid """

        self.m_grid = grid

    def set_sfunc_and_quads(self, basis, quads):
        """ set shape functions and quadrature points """

        (self.m_basis, self.m_quads) = (basis, quads)

        self.m_reac = self.m_basis.get_mass_mat()
        self.m_conv = self.m_basis.get_grad_disp()
        self.m_diff = self.m_basis.get_grad_grad()
        self.m_zero = zeros(self.m_reac.shape)
        self._shape_at_quad_pts()

    def set_cos(self, diff, conv, reac, d_is_real, c_is_real, r_is_real):
        """ set pde coefficients """

        if [] == diff:
            self.m_d = None
        else:
            self.m_d = diff

        if [] == conv:
            self.m_c = None
        else:
            self.m_c = conv

        if [] == reac:
            self.m_r = None
        else:
            self.m_r = reac
        (self.m_c_is_real, self.m_d_is_real, self.m_r_is_real) = \
            (d_is_real, c_is_real, r_is_real)

    def _reaccion_on_seg(self, i, delta_x, seg):
        """ get reaction term on grid """

        if self.m_r is None:
            return self.m_zero

        if self.m_r_is_real:
            return delta_x * self.m_r * self.m_reac

        return self.in_prod1(seg, i, self.m_r, self.m_shape_at_quad)

    def _conv_on_seg(self, i, delta_x, seg):
        """ get convection term on grid """
        if self.m_c is None:
            return self.m_zero

        if self.m_c_is_real:
            return self.m_c * self.m_conv

        # TODO the factor '2' shows up because the reference basis is
        # defined in the interval [-1, 1] which is of length 2.

        return 2. * self.in_prod2(seg, i, self.m_c,
                                 self.m_shape_at_quad,
                                 self.m_der_at_quad) / delta_x

    def diffuse_on_seg(self, i, delta_x, seg):
        """ diffusion term on grid """

        if self.m_d is None:
            return self.m_zero

        if self.m_d_is_real:
            return -self.m_d * self.m_diff / delta_x

        return -4 * self.in_prod1(seg, i, self.m_d, self.m_der_at_quad) \
            / delta_x ** 2

    def __mass_matrix(self):
        """
       In finite elements the grid does not have to be uniform
       That's why we have to use loops to assemble the matrices.
    """

        delta_x = self.m_grid.getDx()
        nsegs = self.m_grid.getNumSegs()

        mass_mat = self.m_basis.get_mass_mat()

        dof = self.m_basis.dof()
        dim = (dof - 1) * nsegs + 1
        diag0 = zeros((dim, 1))

        sdiag = diag(mass_mat)
        sdiag = sdiag.reshape((dof, 1))
        for (i, delta_xi) in enumerate(delta_x):
            diag0[(dof - 1) * i:(dof - 1) * i + dof] += sdiag * delta_xi

        diag0 = diag0.flatten()
        if 2 == dof:
            diag1 = mass_mat[0, 1] * delta_x.flatten()
            self.m_mass_diag = zeros((diag0.size, 3))

        if 3 == dof:
            diag1 = (diag(mass_mat, 1) * c_[delta_x, delta_x]).flatten()
            diag2 = mass_mat[0, 2] * c_[delta_x,
                    zeros(delta_x.shape)].flatten()
            self.m_mass_diag = zeros((diag0.size, 5))
            self.m_mass_diag[2:, 3] = diag2[:-1]
            self.m_mass_diag[:-1, 4] = diag2

            diag2 = diag2[0:-1].flatten()

        self.m_mass_diag[:, 0] = diag0
        self.m_mass_diag[1:, 1] = diag1.copy()
        self.m_mass_diag[:-1, 2] = diag1.copy()

    def _react_matrix(self):
        """ construct reaction matrix """

        delta_x = self.m_grid.getDx()
        segs = self.m_grid.getSegments()
        dof = self.m_basis.dof()

        dim = (dof - 1) * self.m_grid.getNumSegs() + 1
        diag0 = zeros((dim, 1))
        diag1 = zeros((dim - 1, 1))

        if 2 == dof:
            self.m_reac_diag = zeros((diag0.size, 3))
            for (i, delta_xi) in enumerate(delta_x):
                mass_mat = self._reaccion_on_seg(i, delta_xi, segs[i])
                diag0[(dof - 1) * i:(dof - 1) * i + dof] += \
                    diag(mass_mat).reshape((2, 1))
                diag1[i] = mass_mat[0, 1]

        if 3 == dof:
            diag2 = zeros((dim - 1, 1))
            self.m_reac_diag = zeros((diag0.size, 5))
            for (i, delta_xi) in enumerate(delta_x):
                mass_mat = self._reaccion_on_seg(i, delta_xi, segs[i])
                diag0[(dof - 1) * i:(dof - 1) * i + dof] += \
                    diag(mass_mat).reshape((3, 1))
                diag1[2 * i:2 * (i + 1)] = diag(mass_mat,
                        1).reshape((2, 1))
                diag2[2 * i] = mass_mat[0, 2]
            diag2 = diag2.flatten()
            self.m_reac_diag[2:, 3] = diag2[:-1]
            self.m_reac_diag[:-1, 4] = diag2

        diag0 = diag0.flatten()
        diag1 = diag1.flatten()
        self.m_reac_diag[:, 0] = diag0
        self.m_reac_diag[1:, 1] = diag1.copy()
        self.m_reac_diag[:-1, 2] = diag1.copy()

    def discretize(self):
        """
      assemble the stiffness matrix:
        d*u_xx + c*u_x + r*u

        d ~ diffusion
        c ~ convection
        r ~ reaction
    """

        delta_x = self.m_grid.getDx()
        segs = self.m_grid.getSegments()
        dof = self.m_basis.dof()

        dim = (dof - 1) * self.m_grid.getNumSegs() + 1
        diag0 = zeros((dim, 1))
        diag1 = zeros((dim - 1, 1))
        diag_1 = zeros((dim - 1, 1))

        if 2 == dof:
            self.m_stiff_diag = zeros((diag0.size, 3))
            for (i, delta_xi) in enumerate(delta_x):
                conv_mat = self._conv_on_seg(i, delta_xi, segs[i])
                diff_mat = self.diffuse_on_seg(i, delta_xi, segs[i])
                tmp = conv_mat + diff_mat
                diag0[(dof - 1) * i:(dof - 1) * i + dof] += \
                    diag(tmp).reshape((2, 1))
                diag1[i] = tmp[0, 1]
                diag_1[i] = tmp[1, 0]

        if 3 == dof:
            diag2 = zeros((dim - 1, 1))
            diag_2 = zeros((dim - 1, 1))
            self.m_stiff_diag = zeros((diag0.size, 5))
            for (i, delta_xi) in enumerate(delta_x):
                conv_mat = self._conv_on_seg(i, delta_xi, segs[i])
                diff_mat = self.diffuse_on_seg(i, delta_xi, segs[i])
                tmp = conv_mat + diff_mat
                diag0[2 * i:2 * i + 3] += diag(tmp).reshape((3, 1))

                diag1[2 * i:2 * (i + 1)] = diag(tmp, 1).reshape((2, 1))
                diag_1[2 * i:2 * (i + 1)] = diag(tmp, -1).reshape((2,
                        1))

                diag2[2 * i] = tmp[0, 2]
                diag_2[2 * i] = tmp[2, 0]

            diag2 = diag2.flatten()
            diag_2 = diag_2.flatten()
            self.m_stiff_diag[2:, 3] = diag2[:-1]
            self.m_stiff_diag[:-1, 4] = diag_2

        diag0 = diag0.flatten()
        diag1 = diag1.flatten()
        diag_1 = diag_1.flatten()
        self.m_stiff_diag[:, 0] = diag0
        self.m_stiff_diag[1:, 1] = diag1.copy()
        self.m_stiff_diag[:-1, 2] = diag_1.copy()

        if self.m_r is None or self.m_r_is_real and 0 == abs(self.m_r):
            return

        self._react_matrix()
        self.m_stiff_diag += self.m_reac_diag

    def get_pdecos(self):
        """ Coeffients used to assemble the Matrices. """

        return (self.m_d, self.m_c, self.m_r)

    def get_mass_mat(self):
        """ get mass matrix. """

        if self.m_mass_diag is None:
            self.__mass_matrix()
        return self.m_mass_diag

    def get_stiffness(self):
        """ get stiffness matrix. """

        return self.m_stiff_diag

    def in_prod1(self, seg, i, spl, shape_at_quads):
        """ inner product on splines. """

        (xleft, xright) = seg
        length = xright - xleft

        lqpoints = xleft + length * self.m_quads.get_qpoints()
        spl_at_quads = spl.value(lqpoints, i)  # Spline values
        spl_at_quads = (spl_at_quads * self.m_quads.get_weights()
                        / 2.).flatten()  # apply quad weights.

        dim1 = shape_at_quads.shape[0]
        inn = zeros((dim1, dim1))
        for (i, gr0) in enumerate(shape_at_quads):
            gr1 = gr0 * spl_at_quads
            for (j, gr2) in enumerate(shape_at_quads):
                inn[i, j] = length * dot(gr1, gr2)
        return inn

    def in_prod2(self, seg, i, spl, shape_at_quads1, shape_at_quads2):
        """ inner product on splines. """

        (xleft, xright) = seg
        length = xright - xleft

        lqpoints = xleft + length * self.m_quads.get_qpoints()
        spl_at_quads = spl.value(lqpoints, i)  # Spline values
        spl_at_quads = (spl_at_quads * self.m_quads.get_weights()
                        / 2.).flatten()  # apply quad weights.

        dim1 = shape_at_quads1.shape[0]
        inn = zeros((dim1, dim1))
        for (i, gr0) in enumerate(shape_at_quads1):
            gr1 = gr0 * spl_at_quads
            for (j, gr2) in enumerate(shape_at_quads2):
                inn[i, j] = length * dot(gr1, gr2)
        return inn

    def _shape_at_quad_pts(self):
        """ values shape function at quad points. """

        qpoints = self.m_quads.get_qpoints()
        #xlimits = self.m_basis.shape_func_support()

        #qpoints1 = (xlimits[1] - xlimits[0]) * (1. + qpoints + xlimits[0]) / 2.
        self.m_der_at_quad = self.m_basis.derivatives(qpoints)
        self.m_shape_at_quad = self.m_basis.values(qpoints)
        assert self.m_der_at_quad.shape == self.m_shape_at_quad.shape


